{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Ochre.Compiler.CompilerMon
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Monad supplying fresh identifiers plus error 
-- reporting.
--
--------------------------------------------------------------------------------


module Ochre.Compiler.CompilerMon
  ( 

    ErrMsg
  , Compiler
  , runCompiler

  , initState
  , fresh
  , refresh
  , throwError

  ) where

import Ochre.Compiler.VarId

import Control.Applicative

type ErrMsg = String



-- | Note - may need extending for to account for Zak \"ports\".
--
data St = St 
    { ivar_count :: !Int
    , kvar_count :: !Int
    , avar_count :: !Int
    } 

state_zero :: St
state_zero = St { ivar_count = 1, kvar_count = 1, avar_count = 1 }
    



newtype Compiler a = CM { getCM :: St -> Either ErrMsg (a,St) }

instance Functor Compiler where
  fmap f ma = CM $ \s -> getCM ma s >>= \(a,s1) -> return (f a, s1)

instance Applicative Compiler where
  pure a    = CM $ \s -> return (a,s)
  af <*> av = CM $ \s -> getCM af s  >>= \(f,s1) ->
                         getCM av s1 >>= \(a,s2) ->
                         return (f a,s2)

instance Monad Compiler where
  return   = pure
  ma >>= k = CM $ \s -> getCM ma s >>= \(a,s1) -> getCM (k a) s1


runCompiler :: Compiler a -> Either ErrMsg a
runCompiler ma = fmap fst $ getCM ma state_zero

initState :: Compiler () 
initState = CM $ \_ -> return ((), state_zero)

getCount :: RateId -> St -> Int
getCount I_RATE = ivar_count
getCount K_RATE = kvar_count
getCount A_RATE = avar_count


incrCount :: RateId -> St -> St
incrCount I_RATE = (\s i -> s { ivar_count = i+1 }) <*> ivar_count
incrCount K_RATE = (\s i -> s { kvar_count = i+1 }) <*> kvar_count
incrCount A_RATE = (\s i -> s { avar_count = i+1 }) <*> avar_count 


fresh :: RateId -> Compiler VarId
fresh ri = CM $ \st -> 
    let i = getCount ri st
    in return (Local ri i, incrCount ri st)



-- | Add a fresh subscript to the variable name.
--
refresh :: VarId -> Compiler VarId
refresh (Local ri _) = CM $ \st -> 
    let i = getCount ri st 
    in return (Local ri i, incrCount ri st)
refresh varid        = CM $ \st -> return (varid, st)



throwError :: ErrMsg -> Compiler a
throwError err = CM $ \_ -> Left err
